package com.ruoyi.ai.factory.impl;

import com.alibaba.dashscope.common.Role;
import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.core.auth.Auth;
import com.baidubce.qianfan.core.builder.ChatBuilder;
import com.ruoyi.ai.domain.vo.WxChatVo;
import com.ruoyi.ai.domain.vo.WxModelConfigVo;
import com.ruoyi.ai.enums.MfgEnum;
import com.ruoyi.ai.factory.ModelChatTemplate;
import com.ruoyi.ai.ws.ChatWebSocketManager;
import com.ruoyi.common.exception.ServiceException;
import com.ruoyi.common.utils.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.List;
import java.util.function.Consumer;

/**
 * 百度千帆
 */
@Slf4j
@Service
public class QFanModelChatService extends ModelChatTemplate {
    private Qianfan qianfan;

    @Resource
    private ChatWebSocketManager chatWebSocketManager;

    public QFanModelChatService() {
        super(MfgEnum.QFAN.getType());
    }

    @Override
    public String chatStream(List<WxChatVo> chatVoList, Long userId) {
        try {
            log.info("发送的请求为{}", chatVoList);

            // 定义消息处理函数，用于发送websocket消息
            Consumer<String> messageHandler = message -> {
                chatWebSocketManager.sendOneMessage(userId, message);
            };

            // 提取系统消息
            String systemPrompt = null;
            for (WxChatVo wxChatVo : chatVoList) {
                if (Role.SYSTEM.getValue().equals(wxChatVo.getRole())) {
                    systemPrompt = wxChatVo.getContent();
                    break;
                }
            }

            // 构建ChatBuilder并添加消息
            ChatBuilder chatBuilder = getChatBuilderBase(systemPrompt);
            for (WxChatVo wxChatVo : chatVoList) {
                // 跳过系统消息，因为已经在 getChatBuilderBase 中设置
                if (!Role.SYSTEM.getValue().equals(wxChatVo.getRole())) {
                    chatBuilder.addMessage(wxChatVo.getRole(), wxChatVo.getContent());
                }
            }

            // 使用公共方法执行流式聊天
            String result = doChatStream(chatBuilder, messageHandler);

            //发送对话结束标记，前端好判断本次对话已结束，允许下一轮对话
            chatWebSocketManager.sendOneMessage(userId, OVER_FLAG);
            return result;
        } catch (Exception e) {
            log.error("调用百度接口出现问题：{}", e.getMessage());
            throw new ServiceException("调用百度接口出现问题");
        }
    }

    @Override
    public String chatStream(List<WxChatVo> chatVoList, Long userId, String prompt) {
        try {
            log.info("发送的请求为{}", chatVoList);

            // 定义消息处理函数，用于发送websocket消息
            Consumer<String> messageHandler = message -> {
                chatWebSocketManager.sendOneMessage(userId, message);
            };

            // 构建ChatBuilder并添加消息
            ChatBuilder chatBuilder = getChatBuilderBase(prompt);
            for (WxChatVo wxChatVo : chatVoList) {
                // 跳过系统消息，因为已经在 getChatBuilderBase 中设置了 prompt
                if (!Role.SYSTEM.getValue().equals(wxChatVo.getRole())) {
                    chatBuilder.addMessage(wxChatVo.getRole(), wxChatVo.getContent());
                }
            }

            // 使用公共方法执行流式聊天
            String result = doChatStream(chatBuilder, messageHandler);

            //发送对话结束标记，前端好判断本次对话已结束，允许下一轮对话
            chatWebSocketManager.sendOneMessage(userId, OVER_FLAG);
            return result;
        } catch (Exception e) {
            log.error("调用百度接口出现问题：{}", e.getMessage());
            throw new ServiceException("调用百度接口出现问题");
        }
    }

    @Override
    public String chatStreamOnce(String content) {
        try {
            // 构建ChatBuilder并添加消息
            ChatBuilder chatBuilder = getChatBuilderBase(null);
            chatBuilder.addMessage(Role.USER.getValue(), content);

            // 使用公共方法执行流式聊天，不发送websocket消息
            return doChatStream(chatBuilder, null);
        } catch (Exception e) {
            log.error("调用百度接口出现问题：{}", e.getMessage());
            return "调用大模型接口失败，请稍后尝试";
        }
    }

    @Override
    public String chatStreamOnce(String content, String prompt) {
        try {
            // 构建ChatBuilder并添加消息，直接传入prompt
            ChatBuilder chatBuilder = getChatBuilderBase(prompt);
            
            // 添加用户消息
            chatBuilder.addMessage(Role.USER.getValue(), content);

            // 使用公共方法执行流式聊天，不发送websocket消息
            return doChatStream(chatBuilder, null);
        } catch (Exception e) {
            log.error("调用百度接口出现问题：{}", e.getMessage());
            return "调用大模型接口失败，请稍后尝试";
        }
    }

    /**
     * 执行流式聊天的公共方法
     * @param chatBuilder 聊天构建器
     * @param messageHandler 消息处理函数，如果为null则不发送websocket消息
     * @return 聊天结果
     */
    private String doChatStream(ChatBuilder chatBuilder, Consumer<String> messageHandler) {
        // 结果拼接对象
        StringBuilder resultBuilder = new StringBuilder();

        chatBuilder.executeStream() // 发起流式请求
            .forEachRemaining(chunk -> {
                String result = chunk.getResult();
                resultBuilder.append(result);

                // 如果有消息处理函数，则发送消息
                if (messageHandler != null) {
                    messageHandler.accept(result);
                }
            }); // 流式迭代，并处理消息

        //将最终结果返回
        return resultBuilder.toString();
    }

    /**
     * 获取基础的ChatBuilder
     * @param prompt 系统提示词，可以为null
     * @return 配置好的ChatBuilder
     */
    public ChatBuilder getChatBuilderBase(String prompt) {
        WxModelConfigVo config = getConfig();
        if (config == null) {
            throw new ServiceException("模型参数未配置");
        }
        if (qianfan == null) {
            //此为旧版应用的ak,sk鉴权
            //qianfan = new Qianfan(Auth.TYPE_OAUTH, config.getApiKey(), config.getSecretKey());
            //此为新版的IAM鉴权，更简单方便了，不用再申请应用的ak,sk了
            //百度智能云IAM地址：https://console.bce.baidu.com/iam/#/iam/accesslist
            qianfan = new Qianfan(config.getApiKey(),config.getSecretKey());
        }
        ChatBuilder chatBuilder = qianfan.chatCompletion();
        chatBuilder.model(config.getModelType());
        
        // 设置系统提示词
        if (StringUtils.isNotBlank(prompt)) {
            chatBuilder.system(prompt);
        }
        
        return chatBuilder;
    }
}
